{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f6dc6613-b498-42f1-80f7-0b3633c269c6",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'tushare'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtushare\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mts\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'tushare'"
     ]
    }
   ],
   "source": [
    "import tushare as ts\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.metrics import accuracy_score\n",
    "import time\n",
    "import os\n",
    "# 设置tushare token\n",
    "ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')\n",
    "pro = ts.pro_api()\n",
    "class QuantTradingSystem:\n",
    "    def __init__(self, start_date, end_date, stock_list, initial_capital=1000000):\n",
    "        self.start_date = start_date\n",
    "        self.end_date = end_date\n",
    "        self.stock_list = stock_list\n",
    "        self.initial_capital = initial_capital\n",
    "        self.trading_results = {}\n",
    "    def get_stock_data(self, ts_code):\n",
    "        \"\"\"获取股票数据和财务指标\"\"\"\n",
    "        # 获取日线行情\n",
    "        df_daily = ts.pro_bar(ts_code=ts_code, adj='qfq', \n",
    "                              start_date=self.start_date, end_date=self.end_date)\n",
    "        if df_daily is None or len(df_daily) == 0:\n",
    "            return None\n",
    "        df_daily = df_daily.sort_values('trade_date')\n",
    "        # 获取财务指标\n",
    "        try:\n",
    "            # 获取季度财务指标\n",
    "            df_finance = pro.fina_indicator(ts_code=ts_code, \n",
    "                                           start_date=self.start_date[:4]+'0101', \n",
    "                                           end_date=self.end_date[:4]+'1231')\n",
    "            if df_finance is not None and len(df_finance) > 0:\n",
    "                df_finance = df_finance.sort_values('end_date')\n",
    "                # 财务数据映射到每个交易日\n",
    "                df_daily['end_date'] = df_daily['trade_date'].apply(lambda x: x[:6]+'01')\n",
    "                df_merged = pd.merge(df_daily, df_finance, on=['ts_code', 'end_date'], how='left')\n",
    "                # 使用前向填充处理缺失值\n",
    "                df_merged = df_merged.ffill()\n",
    "                return df_merged\n",
    "            else:\n",
    "                return df_daily\n",
    "        except Exception as e:\n",
    "            print(f\"获取{ts_code}财务数据出错: {e}\")\n",
    "            return df_daily\n",
    "    def calculate_technical_indicators(self, df):\n",
    "        \"\"\"计算技术指标\"\"\"\n",
    "        if df is None:\n",
    "            return None\n",
    "        # 计算移动平均线\n",
    "        df['MA5'] = df['close'].rolling(window=5).mean()\n",
    "        df['MA10'] = df['close'].rolling(window=10).mean()\n",
    "        df['MA20'] = df['close'].rolling(window=20).mean()\n",
    "        # 计算相对强弱指标(RSI)\n",
    "        delta = df['close'].diff()\n",
    "        gain = delta.where(delta > 0, 0)\n",
    "        loss = -delta.where(delta < 0, 0)\n",
    "        avg_gain = gain.rolling(window=14).mean()\n",
    "        avg_loss = loss.rolling(window=14).mean()\n",
    "        rs = avg_gain / avg_loss\n",
    "        df['RSI'] = 100 - (100 / (1 + rs))\n",
    "        # 计算MACD\n",
    "        df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()\n",
    "        df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()\n",
    "        df['MACD'] = df['EMA12'] - df['EMA26']\n",
    "        df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()\n",
    "        df['Histogram'] = df['MACD'] - df['Signal']\n",
    "        # 计算波动率\n",
    "        df['Volatility'] = df['close'].rolling(window=20).std()\n",
    "        # 计算每日收益率\n",
    "        df['daily_return'] = df['close'].pct_change()\n",
    "        # 计算未来收益率(用于标签)\n",
    "        df['future_return'] = df['close'].pct_change(5).shift(-5)\n",
    "        return df\n",
    "    def financial_strategy(self, df):\n",
    "        \"\"\"基于财务指标的交易策略\"\"\"\n",
    "        if df is None or len(df) == 0:\n",
    "            return None  \n",
    "        signals = pd.DataFrame(index=df.index)\n",
    "        signals['signal'] = 0\n",
    "        # 使用财务指标生成信号\n",
    "        # ROE(净资产收益率) > 15% 且毛利率 > 30% 且净利润增长率 > 10%\n",
    "        if 'roe' in df.columns and 'grossprofit_margin' in df.columns and 'profit_to_gr' in df.columns:\n",
    "            buy_condition = (df['roe'] > 15) & (df['grossprofit_margin'] > 30) & (df['profit_to_gr'] > 10)\n",
    "            # 结合技术指标确认\n",
    "            ma_condition = (df['MA5'] > df['MA20']) & (df['MA5'] > df['MA10'])\n",
    "            rsi_condition = df['RSI'] < 70\n",
    "            signals.loc[buy_condition & ma_condition & rsi_condition, 'signal'] = 1 \n",
    "            # 卖出信号：ROE下降 或 毛利率下降 或 RSI>80 或 价格跌破MA20\n",
    "            sell_condition = ((df['roe'] < df['roe'].shift(1)) | \n",
    "                             (df['grossprofit_margin'] < df['grossprofit_margin'].shift(1)) |\n",
    "                             (df['RSI'] > 80) |\n",
    "                             (df['close'] < df['MA20']))\n",
    "            signals.loc[sell_condition, 'signal'] = -1 \n",
    "        return signals\n",
    "    def ml_strategy(self, df):\n",
    "        \"\"\"基于机器学习的交易策略\"\"\"\n",
    "        if df is None or len(df) < 100:  # 至少需要100个样本\n",
    "            return None \n",
    "        # 准备特征\n",
    "        features = ['MA5', 'MA10', 'MA20', 'RSI', 'MACD', 'Signal', 'Histogram', 'Volatility']\n",
    "        if 'roe' in df.columns:\n",
    "            features.extend(['roe', 'grossprofit_margin', 'profit_to_gr', 'debt_to_assets'])\n",
    "        # 移除包含NaN的行\n",
    "        df_clean = df.dropna(subset=features + ['future_return'])\n",
    "        if len(df_clean) < 50:  # 确保有足够的样本\n",
    "            return None\n",
    "        # 创建标签：1表示上涨，0表示下跌\n",
    "        df_clean['label'] = np.where(df_clean['future_return'] > 0.01, 1, 0)\n",
    "        # 划分训练集和测试集\n",
    "        X = df_clean[features]\n",
    "        y = df_clean['label']\n",
    "        # 使用前80%的数据作为训练集，后20%作为测试集\n",
    "        split_idx = int(len(df_clean) * 0.8)\n",
    "        X_train, X_test = X.iloc[:split_idx], X.iloc[split_idx:]\n",
    "        y_train, y_test = y.iloc[:split_idx], y.iloc[split_idx:]\n",
    "        # 训练随机森林分类器\n",
    "        model = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "        model.fit(X_train, y_train)\n",
    "        # 在测试集上的预测\n",
    "        y_pred = model.predict(X_test)\n",
    "        test_accuracy = accuracy_score(y_test, y_pred)\n",
    "        print(f\"模型准确率: {test_accuracy:.4f}\")\n",
    "        # 生成交易信号\n",
    "        signals = pd.DataFrame(index=df.index)\n",
    "        signals['signal'] = 0\n",
    "        # 将预测结果映射到对应的日期\n",
    "        test_dates = df_clean.index[split_idx:]\n",
    "        signals.loc[test_dates, 'signal'] = y_pred\n",
    "        # 只在预测上涨时买入，卖出信号使用技术指标\n",
    "        sell_condition = (df['RSI'] > 80) | (df['close'] < df['MA20'])\n",
    "        signals.loc[sell_condition, 'signal'] = -1\n",
    "        return signals\n",
    "    def backtest(self, df, signals, strategy_name, ts_code):\n",
    "        \"\"\"回测交易策略\"\"\"\n",
    "        if df is None or signals is None:\n",
    "            return None\n",
    "        # 合并信号和价格数据\n",
    "        df_with_signals = pd.merge(df, signals, left_index=True, right_index=True, how='left')\n",
    "        df_with_signals['signal'] = df_with_signals['signal'].fillna(0)\n",
    "        # 初始化持仓和现金\n",
    "        positions = pd.DataFrame(index=df_with_signals.index).fillna(0.0)\n",
    "        positions[ts_code] = 100 * df_with_signals['signal']  # 假设每次交易100股\n",
    "        # 计算投资组合价值\n",
    "        portfolio = positions.multiply(df_with_signals['close'], axis=0)\n",
    "        pos_diff = positions.diff()  # 交易变化\n",
    "        # 计算现金余额\n",
    "        portfolio['holdings'] = (positions.multiply(df_with_signals['close'], axis=0)).sum(axis=1)\n",
    "        portfolio['cash'] = self.initial_capital - (pos_diff.multiply(df_with_signals['close'], axis=0)).sum(axis=1).cumsum()\n",
    "        portfolio['total'] = portfolio['cash'] + portfolio['holdings']\n",
    "        portfolio['returns'] = portfolio['total'].pct_change()\n",
    "        # 计算策略评价指标\n",
    "        cumulative_return = (portfolio['total'].iloc[-1] / portfolio['total'].iloc[0]) - 1\n",
    "        sharpe_ratio = np.sqrt(252) * (portfolio['returns'].mean() / portfolio['returns'].std()) if portfolio['returns'].std() != 0 else 0\n",
    "        # 计算最大回撤\n",
    "        portfolio['cumulative'] = (1 + portfolio['returns']).cumprod()\n",
    "        portfolio['cumulative_max'] = portfolio['cumulative'].cummax()\n",
    "        portfolio['drawdown'] = (portfolio['cumulative'] / portfolio['cumulative_max']) - 1\n",
    "        max_drawdown = portfolio['drawdown'].min()\n",
    "        results = {\n",
    "            'cumulative_return': cumulative_return,\n",
    "            'sharpe_ratio': sharpe_ratio,\n",
    "            'max_drawdown': max_drawdown,\n",
    "            'portfolio': portfolio,\n",
    "            'df_with_signals': df_with_signals\n",
    "        }\n",
    "        return results\n",
    "    def run_strategy(self, strategy_func, strategy_name):\n",
    "        \"\"\"运行策略并汇总结果\"\"\"\n",
    "        all_results = {}\n",
    "        for ts_code in self.stock_list:\n",
    "            print(f\"回测 {ts_code} {strategy_name} 策略...\")\n",
    "            # 获取数据\n",
    "            df = self.get_stock_data(ts_code)\n",
    "            if df is None:\n",
    "                print(f\"无法获取 {ts_code} 的数据\")\n",
    "                continue   \n",
    "            # 计算技术指标\n",
    "            df = self.calculate_technical_indicators(df)  \n",
    "            # 生成交易信号\n",
    "            signals = strategy_func(df)\n",
    "            # 回测策略\n",
    "            results = self.backtest(df, signals, strategy_name, ts_code)\n",
    "            if results:\n",
    "                all_results[ts_code] = results\n",
    "                print(f\"{ts_code} {strategy_name} 策略回测完成\")\n",
    "                print(f\"累计收益率: {results['cumulative_return']:.4f}\")\n",
    "                print(f\"夏普比率: {results['sharpe_ratio']:.4f}\")\n",
    "                print(f\"最大回撤: {results['max_drawdown']:.4f}\")\n",
    "                print(\"-\" * 50)\n",
    "            # 避免API调用频率限制\n",
    "            time.sleep(0.5)\n",
    "        self.trading_results[strategy_name] = all_results\n",
    "        return all_results\n",
    "    def compare_strategies(self):\n",
    "        \"\"\"比较不同策略的表现\"\"\"\n",
    "        strategies = list(self.trading_results.keys())\n",
    "        if len(strategies) < 2:\n",
    "            print(\"至少需要运行两种策略才能进行比较\")\n",
    "            return  \n",
    "        # 创建结果表格\n",
    "        comparison = pd.DataFrame(columns=['股票代码', '财务策略收益率', '机器学习策略收益率', '最优策略'])\n",
    "        for ts_code in self.stock_list:\n",
    "            if all(ts_code in self.trading_results[strategy] for strategy in strategies):\n",
    "                returns = {}\n",
    "                for strategy in strategies:\n",
    "                    returns[strategy] = self.trading_results[strategy][ts_code]['cumulative_return']     \n",
    "                best_strategy = max(returns, key=returns.get)  \n",
    "                comparison = pd.concat([comparison, pd.DataFrame({\n",
    "                    '股票代码': [ts_code],\n",
    "                    '财务策略收益率': [returns.get('financial', float('nan'))],\n",
    "                    '机器学习策略收益率': [returns.get('ml', float('nan'))],\n",
    "                    '最优策略': [best_strategy]\n",
    "                })])\n",
    "        return comparison\n",
    "    def plot_results(self, ts_code, strategy_names):\n",
    "        \"\"\"绘制策略表现图表\"\"\"\n",
    "        plt.figure(figsize=(14, 10))\n",
    "        for i, strategy_name in enumerate(strategy_names):\n",
    "            if ts_code in self.trading_results[strategy_name]:\n",
    "                portfolio = self.trading_results[strategy_name][ts_code]['portfolio']\n",
    "     \n",
    "                plt.subplot(len(strategy_names), 1, i+1)\n",
    "                plt.plot(portfolio['total'])\n",
    "                plt.title(f\"{ts_code} {strategy_name} 策略累计收益\")\n",
    "                plt.xlabel('日期')\n",
    "                plt.ylabel('资产总值')\n",
    "                plt.grid(True)           \n",
    "                # 标注交易点\n",
    "                df_with_signals = self.trading_results[strategy_name][ts_code]['df_with_signals']\n",
    "                buy_points = df_with_signals[df_with_signals['signal'] == 1]\n",
    "                sell_points = df_with_signals[df_with_signals['signal'] == -1]               \n",
    "                for date, row in buy_points.iterrows():\n",
    "                    plt.axvline(x=date, color='g', linestyle='--', alpha=0.5)   \n",
    "                for date, row in sell_points.iterrows():\n",
    "                    plt.axvline(x=date, color='r', linestyle='--', alpha=0.5)\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(f\"{ts_code}_strategy_comparison.png\")\n",
    "        plt.close()\n",
    "# 主程序\n",
    "if __name__ == \"__main__\":\n",
    "    # 回测参数设置\n",
    "    start_date = '20200101'\n",
    "    end_date = '20231231'\n",
    "    # 选择要回测的股票池\n",
    "    stock_list = [\n",
    "        '000001.SZ',  # 平安银行\n",
    "        '000858.SZ',  # 五粮液\n",
    "        '002594.SZ',  # 比亚迪\n",
    "        '600519.SH',  # 贵州茅台\n",
    "        '601318.SH',  # 中国平安\n",
    "        '601899.SH',  # 紫金矿业\n",
    "        '603259.SH',  # 药明康德\n",
    "        '603986.SH',  # 兆易创新\n",
    "        '688122.SH',  # 西部超导\n",
    "        '300750.SZ',  # 宁德时代\n",
    "        '600036.SH',  # 招商银行\n",
    "        '600276.SH',  # 恒瑞医药\n",
    "        '600900.SH',  # 长江电力\n",
    "        '601012.SH',  # 隆基绿能\n",
    "        '601688.SH',  # 华泰证券\n",
    "        '601888.SH',  # 中国中免\n",
    "        '603501.SH',  # 韦尔股份\n",
    "        '688981.SH',  # 中芯国际\n",
    "        '002415.SZ',  # 海康威视\n",
    "        '002714.SZ',  # 牧原股份\n",
    "    ]\n",
    "    # 初始化量化交易系统\n",
    "    qts = QuantTradingSystem(start_date, end_date, stock_list)\n",
    "    # 运行财务指标策略\n",
    "    financial_results = qts.run_strategy(qts.financial_strategy, 'financial')\n",
    "    # 运行机器学习策略\n",
    "    ml_results = qts.run_strategy(qts.ml_strategy, 'ml')\n",
    "    # 比较两种策略\n",
    "    comparison = qts.compare_strategies()\n",
    "    print(\"\\n策略比较结果:\")\n",
    "    print(comparison)\n",
    "    # 绘制策略表现图表\n",
    "    for ts_code in stock_list:\n",
    "        qts.plot_results(ts_code, ['financial', 'ml'])\n",
    "    # 计算总体表现\n",
    "    print(\"\\n总体表现统计:\")\n",
    "    for strategy in ['financial', 'ml']:\n",
    "        total_return = 0\n",
    "        positive_count = 0\n",
    "        for ts_code, result in qts.trading_results[strategy].items():\n",
    "            if result['cumulative_return'] > 0:\n",
    "                positive_count += 1\n",
    "            total_return += result['cumulative_return']\n",
    "        \n",
    "        avg_return = total_return / len(qts.trading_results[strategy]) if qts.trading_results[strategy] else 0\n",
    "        win_rate = positive_count / len(qts.trading_results[strategy]) if qts.trading_results[strategy] else 0       \n",
    "        print(f\"{strategy} 策略:\")\n",
    "        print(f\"平均收益率: {avg_return:.4f}\")\n",
    "        print(f\"胜率: {win_rate:.2%}\")\n",
    "        print(\"-\" * 30)    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:base] *",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
